using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.IO;
using System.Linq;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Content.Pipeline;
using Microsoft.Xna.Framework.Content.Pipeline.Graphics;
using Microsoft.Xna.Framework.Content.Pipeline.Processors;
using Oops.Xna.Framework.Graphics;

namespace Oops.Xna.Framework.Content.Pipeline
{
	public class AnimationDefinition
	{
		public string Name { get; set; }
		public TimeSpan StartTime { get; set; }
		public TimeSpan EndTime { get; set; }
	}

	[ContentProcessor(DisplayName = "Skinned Model - Oops! Framework")]
	public class SkinnedModelProcessor : ModelProcessor
	{
		private const int MaximumBoneCount = 49;

		private List<Matrix> bindPose = new List<Matrix>();
		private List<Matrix> inverseBindPose = new List<Matrix>();
		private List<int> skeletonHierarchy = new List<int>();

		private int sampleRate = 60;

		private Texture2DContent animationTexture;

		/// <summary>
		/// Gets or sets the sample rate of any animation associated with the asset being processed.
		/// </summary>
		[DisplayName("Sample Rate")]
		public int SampleRate
		{
			get { return this.sampleRate; }
			set { this.sampleRate = value; }
		}

		[DisplayName("Animation Definition File Name")]
		public string AnimationDefinitionFileName { get; set; }
		
		#region ModelProcessor Members

		public override ModelContent Process(NodeContent input, ContentProcessorContext context)
		{
			{ ///////////////////////////////////////////////////////////////
				if (input is BoneContent)
				{
					NodeContent node = new NodeContent();
					node.Name = input.Name;
					node.Transform = input.Transform;
					for (int i = input.Children.Count - 1; i >= 0; i--)
					{
						NodeContent child = input.Children[i];

						input.Children.Remove(child);
						node.Children.Insert(0, child);
					}

					foreach (KeyValuePair<string, AnimationContent> channel in input.Animations)
					{
						node.Children[1].Animations.Add(channel.Key, channel.Value);
					}
					input.Animations.Clear();
					input = node;
				}
			} ///////////////////////////////////////////////////////////////

			// Find all skinless meshes.
			var skinlessMeshes = from mesh in input.Children
									 .SelectRecursive(s => s.Children)
									 .Union(new NodeContent[] { input })
									 .OfType<MeshContent>()
								 where mesh.Geometry.All(g => !g.Vertices.Channels.Contains(VertexChannelNames.Weights(0)))
								 select mesh;

			// Remove skinless meshes.
			foreach (var mesh in skinlessMeshes.Reverse<MeshContent>())
			{
				mesh.Parent.Children.Remove(mesh);
			}
			
			// Find all geometry with no weight information.
			var weightlessGeometry = from mesh in input.Children
										 .SelectRecursive(s => s.Children)
										 .Union(new NodeContent[] { input })
										 .OfType<MeshContent>()
									 from geometry in mesh.Geometry.Where(g => g.Vertices.Channels.Contains(VertexChannelNames.Weights(0)))
									 let weightChannel = from weights in geometry.Vertices.Channels.Get<BoneWeightCollection>(VertexChannelNames.Weights(0)) select weights
									 where weightChannel.Any(w => w.Count == 0)
									 select geometry;

			// Remove all weightless geometry.
			foreach (var geometry in weightlessGeometry)
			{
				geometry.Parent.Geometry.Remove(geometry);
			}
			
			// Find all skinned bone names.
			var skinnedBoneNames = (from mesh in input.Children
										.SelectRecursive(s => s.Children)
										.Union(new NodeContent[] { input })
										.OfType<MeshContent>()
									from geometry in mesh.Geometry.Where(g => g.Vertices.Channels.Contains(VertexChannelNames.Weights(0)))
									from weights in geometry.Vertices.Channels.Get<BoneWeightCollection>(VertexChannelNames.Weights(0))
									from weight in weights
									where weight.Weight > 0.0f
									select weight.BoneName).Distinct().ToList();

			// Find the skeleton.
			BoneContent skeleton = MeshHelper.FindSkeleton(input);
			if (skeleton == null)
			{
				throw new InvalidContentException("Input skeleton not found.");
			}

			// Transform nodes by parameters.
			if (this.RotationX != 0.0f || 
				this.RotationY != 0.0f || 
				this.RotationZ != 0.0f || 
				this.Scale != 1.0f)
			{
				Matrix rotationZ = Matrix.CreateRotationZ(MathHelper.ToRadians(this.RotationZ));
				Matrix rotationX = Matrix.CreateRotationX(MathHelper.ToRadians(this.RotationX));
				Matrix rotationY = Matrix.CreateRotationY(MathHelper.ToRadians(this.RotationY));

				Matrix transform = rotationZ * rotationX * rotationY * Matrix.CreateScale(this.Scale);
				MeshHelper.TransformScene(input, transform);
			}

			// Bake the transforms.
			input.Children
				.Where(c => c != skeleton)
				.SelectRecursive(s => s.Children.Where(c => c != skeleton))
				.ToList()
				.ForEach(s => { MeshHelper.TransformScene(s, s.Transform); s.Transform = Matrix.Identity; });

			// Find all skinless bones.
			var skinlessBones = (from bone in input.Children
									 .SelectRecursive(s => s.Children)
									 .OfType<BoneContent>()
								 where !skinnedBoneNames.Contains(bone.Name) && !bone.Name.StartsWith("ref_") && bone != skeleton
								 select bone).ToList();

			// Remove bones that aren't uses for skinning.  Don't remove reference bones.
			for (int i = skinlessBones.Count - 1; i >= 0; i--)
			{
				NodeContent parent = skinlessBones[i].Parent;
				NodeContentCollection children = skinlessBones[i].Children;

				if (parent != null && children.Count == 0)
				{
					parent.Children.Remove(skinlessBones[i]);

					//for (int j = children.Count - 1; j >= 0; j--)
					//{
					//    NodeContent child = children[j];

					//    skinlessBones[i].Children.Remove(child);
					//    parent.Children.Add(child);
					//}
				}
			}

			// Read the bind pose and skeleton hierarchy data.
			var bones = MeshHelper.FlattenSkeleton(skeleton);
			if (bones.Count > SkinnedModelProcessor.MaximumBoneCount)
			{
				throw new InvalidContentException(string.Format("The skeleton has {0} bones, but the maximum supported is for {1}.", bones.Count, SkinnedModelProcessor.MaximumBoneCount));
			}

			// Get bone information (similar to xna example).
			Dictionary<string, int> boneIndices = new Dictionary<string, int>();
			foreach (var bone in bones)
			{
				this.bindPose.Add(bone.Transform);
				this.inverseBindPose.Add(Matrix.Invert(bone.AbsoluteTransform));
				this.skeletonHierarchy.Add(bones.IndexOf(bone.Parent as BoneContent));
				boneIndices.Add(bone.Name, boneIndices.Count);
			}

			// Get the animations and separate them if necessary.
			if (!string.IsNullOrEmpty(this.AnimationDefinitionFileName))
			{
				var animationDefinitions =
					context.BuildAndLoadAsset<object, object>(
						new ExternalReference<object>(Path.GetFullPath(this.AnimationDefinitionFileName)), 
						"PassThroughProcessor") as AnimationDefinition[];

				if (skeleton.Animations.Count > 0)
				{
					var oldAnimationContent = skeleton.Animations.FirstOrDefault();
					skeleton.Animations.Remove(oldAnimationContent.Key);

					foreach (var definition in animationDefinitions)
					{
						TimeSpan startTime = definition.StartTime;
						TimeSpan endTime = definition.EndTime;

						var animationChannels = oldAnimationContent.Value.Channels;
						if (animationChannels.Count() > 0)
						{
							AnimationContent newAnimationContent = new AnimationContent();
							newAnimationContent.Name = definition.Name;
							newAnimationContent.Duration = endTime - startTime;

							foreach (var channel in animationChannels)
							{
								var keyframes = from item in channel.Value
												where item.Time >= startTime && item.Time <= endTime
												select item;

								AnimationChannel newAnimationChannel = new AnimationChannel();
								bool first = true;
								foreach (var keyframe in keyframes)
								{
									newAnimationChannel.Add(
										new AnimationKeyframe(first ? TimeSpan.Zero : keyframe.Time.Subtract(startTime), keyframe.Transform));

									first = false;
								}
								newAnimationContent.Channels.Add(channel.Key, newAnimationChannel);
							}

							skeleton.Animations.Add(newAnimationContent.Name, newAnimationContent);
						}
					}
				}
			}

			// Get the animation clips and all key frame matrices.
			List<Matrix[]> keyFrameMatrices = new List<Matrix[]>();
			IList<AnimationClip> animationClips = this.ProcessAnimations(
				skeleton.Animations, 
				bones, 
				keyFrameMatrices);

			// Create the animation texture.
			var animationContent = SkinnedModelProcessor.GetEncodedTexture(keyFrameMatrices, bones.Count);
			this.animationTexture = new Texture2DContent();
			this.animationTexture.Faces[0].Add(animationContent);

			// The nodes have already been transformed.
			this.RotationX = this.RotationY = this.RotationZ = 0.0f;
			this.Scale = 1.0f;
			
			// Call inherited method.
			ModelContent model = base.Process(input, context);
			
			// Attach the skinning content.
			model.Tag = new SkinningDataContent(
				this.bindPose,
				this.inverseBindPose,
				this.skeletonHierarchy,
				animationTexture,
				animationClips,
				boneIndices);

			// Result.
			return model;
		}

		protected override MaterialContent ConvertMaterial(MaterialContent material, ContentProcessorContext context)
		{
			// Ensure that the material of the model is for a basic effect.
			BasicMaterialContent basicMaterial = material as BasicMaterialContent;
			if (basicMaterial == null)
			{
				throw new InvalidContentException(
					string.Format(
						"SkinnedModelProcessor only supports BasicMaterialContent, but input mesh uses {0}.",
						material.GetType()));
			}

			// Create a new material.
			EffectMaterialContent result = new EffectMaterialContent();
			result.Effect = new ExternalReference<EffectContent>(Path.GetFullPath("Effects\\CommonEffect.fx"));

			// Transfer the diffuse texture.
			if (basicMaterial.Texture != null)
			{
				result.Textures.Add("DiffuseTexture", basicMaterial.Texture); 
			}

			// Create effect parameters for the animation texture deltas.
			if (true)
			{
				result.OpaqueData.Add("BoneDelta", 1.0f / this.animationTexture.Mipmaps[0].Width);
				result.OpaqueData.Add("RowDelta", 1.0f / this.animationTexture.Mipmaps[0].Height);
			}

			// Call inherited method.
			return base.ConvertMaterial(result, context);
		}

		#endregion
		
		private IList<AnimationClip> ProcessAnimations(
			AnimationContentDictionary animations,
			IList<BoneContent> bones,
			List<Matrix[]> outKeyFrames)
		{
			// Create a mapping from the name of the bone to the index.
			var boneMappings = new Dictionary<string, int>();
			for (int i = 0; i < bones.Count; i++)
			{
				string name = bones[i].Name;
				if (!string.IsNullOrEmpty(name))
				{
					boneMappings.Add(name, i);
				}
			}

			// Iterate on all animation content.
			var result = new List<AnimationClip>();
			foreach (var entry in animations)
			{
				// Get the current name and animation content.
				string name = entry.Key;
				AnimationContent animation = entry.Value;

				// Create an animation clip from the the animation content.
				AnimationClip clip = this.ProcessAnimation(
					name, 
					animation, 
					boneMappings,
					outKeyFrames);
				result.Add(clip);
			}

			// Return the animation clips.
			return result;
		}

		private AnimationClip ProcessAnimation(
			string name,
			AnimationContent animation,
			Dictionary<string, int> boneMappings,
			List<Matrix[]> animationRows)
		{
			// Get the key frames for the current animation content.
			var keyframes = SkinnedModelProcessor.GetKeyFrameList(animation, boneMappings);

			// Determine the starting row in the animation texture for
			// the current animation.
			int startRow = animationRows.Count;

			// Calculate the time per frame.
			float timeDelta = 1.0f / (float)this.sampleRate;

			float time = 0.0f;
			int currentKeyFrame = 0;
			Matrix[] boneTransforms = this.bindPose.ToArray();
			while (time < animation.Duration.TotalSeconds)
			{
				for (int i = currentKeyFrame; i < keyframes.Count; i++)
				{
					Keyframe keyframe = keyframes[i];
					if (time < keyframe.Time.TotalSeconds)
					{
						break;
					}

					boneTransforms[keyframe.BoneIndex] = keyframe.Transform;
					currentKeyFrame = i;
				}

				// Determine the world transform of all the bones.
				Matrix[] worldTransforms = new Matrix[boneTransforms.Length];
				worldTransforms[0] = boneTransforms[0];
				for (int i = 1; i < worldTransforms.Length; i++)
				{
					int parentIndex = this.skeletonHierarchy[i];
					worldTransforms[i] = boneTransforms[i] * worldTransforms[parentIndex];
				}

				// Determine the skinning transform of all the bones.
				Matrix[] skinningTransforms = new Matrix[worldTransforms.Length];
				for (int i = 0; i < skinningTransforms.Length; i++)
				{
					skinningTransforms[i] = this.inverseBindPose[i] * worldTransforms[i];
				}

				// Add the skinning transforms to this animation row.
				animationRows.Add(skinningTransforms);

				time += timeDelta;
			}

			// Return the new animation clip created from the animation content.
			return new AnimationClip(
				name, 
				animation.Duration, 
				this.sampleRate, 
				keyframes, 
				startRow, 
				animationRows.Count - 1);
		}

		private static PixelBitmapContent<Vector4> GetEncodedTexture(
			List<Matrix[]> keyFrameMatrices, 
			int boneCount)
		{
			// Create the texture with enough room for the bone times the number
			// of keyframes of al animations.
			var result = new PixelBitmapContent<Vector4>(
				boneCount * 4, 
				keyFrameMatrices.Count);

			// Iterate on the poses of all the animations.
			int y = 0;
			foreach (Matrix[] animationRow in keyFrameMatrices)
			{
				// Iterate on the bone transform of the current pose.
				int x = 0;
				foreach (Matrix transform in animationRow)
				{
					result.SetPixel(x + 0, y, new Vector4(transform.M11, transform.M12, transform.M13, transform.M41));
					result.SetPixel(x + 1, y, new Vector4(transform.M21, transform.M22, transform.M23, transform.M42));
					result.SetPixel(x + 2, y, new Vector4(transform.M31, transform.M32, transform.M33, transform.M43));

					// Increment the bone index.
					x += 4;
				}

				// Iterate the pose index.
				y++;
			}

			// Return the animation texture.
			return result;
		}

		private static List<Keyframe> GetKeyFrameList(
			AnimationContent animation, 
			Dictionary<string, int> boneMappings)
		{
			// Ensure that the animation has a duration set.
			if (animation.Duration <= TimeSpan.Zero)
			{
				throw new InvalidContentException("The animation has a zero duration.");
			}

			// Iterate on all the channels of the animation.
			var result = new List<Keyframe>();
			foreach (var channel in animation.Channels)
			{
				// Get the index of the channel's associated bone.
				int boneIndex;
				if (!boneMappings.TryGetValue(channel.Key, out boneIndex))
				{
					continue;
				}

				// Iterate on all keyframes of the current channel.
				foreach (var keyframe in channel.Value)
				{
					// Create a animation keyframe from the channel's keyframe.
					result.Add(new Keyframe(boneIndex, keyframe.Time, keyframe.Transform));
				}
			}

			// Ensure that there are keyframes for the animation.
			if (result.Count == 0)
			{
				throw new InvalidContentException("The animation has no keyframes.");
			}

			// Sort and return the keyframes.
			result.Sort((a, b) => a.Time.CompareTo(b.Time));
			return result;
		}
	}
}
